

import sys
import numpy as np
from numpy.fft import rfft as fft
from numpy.fft import irfft as ifft
import scipy.signal as signal
import soundfile as sf
# from frequency_domain_adaptive_filters.pfdkf import PFDKF
# from frequency_domain_adaptive_filters.fdkf import FDKF
# from frequency_domain_adaptive_filters.ssaec import SSAEC
from pfdkf import PFDKF

def RMS(x):
  return 10 * np.log10(np.average(x**2) + 1e-6)

def RMS2(x):
  return np.average(x**2)

def dB(x):
  return 10 * np.log10(x + 1e-6)


class ring_buffer:
    def __init__(self, size, delay):
        self.buffer = np.zeros(size)
        self.size = size
        self.head = 0
        self.tail = (size + delay) % size

    def get(self, length):
        if (self.tail + length) <= self.size:
            block = self.buffer[self.tail:self.tail+length]
        else:
            block = np.concatenate((self.buffer[self.tail:], self.buffer[:self.tail+length-self.size]))
        self.tail = (self.tail + length) % self.size
        return block

    def put(self, data):
        length = len(data)
        if (self.head + length) <= self.size:
            self.buffer[self.head:self.head+length] = data
        else:
            self.buffer[self.head:] = data[:self.size - self.head]
            self.buffer[:self.head+length-self.size] = data[self.size - self.head:]
        self.head = (self.head + length) % self.size


tulip_ir = [3.2063674e-06, -1.0982357e-06, 1.082732e-06, -4.6957075e-07, 1.5635995e-07, -2.6467571e-06, 3.1231762e-06, 2.1800389e-07, 5.884862e-06, -6.354895e-06, 1.4848641e-05, -2.8473714e-06, -0.00015431145, -2.0637679e-05, 0.00022234378, 0.000106424726, -0.00015489636, -8.870107e-05, 7.4506104e-05, 7.608519e-05, -8.783229e-06, -4.5082877e-05, -1.7606055e-05, 5.9777562e-06, -1.0221425e-05, 2.153114e-05, 6.5563686e-06, -4.6148463e-05, -2.7593753e-06, 2.1937678e-05, 1.0153939e-05, -1.2419459e-05, -3.1686446e-05, 1.0736528e-05, 3.2484164e-05, -1.1649979e-05, -1.2017728e-05, 1.8915755e-06, 2.0891395e-05, 2.3852415e-06, -8.460686e-06, -2.4144704e-06, 1.6790513e-05, 4.439218e-07, -1.7414071e-05, -8.432164e-06, 4.203739e-06, 4.223816e-06, -4.7690455e-06, -1.6287934e-06, 5.169887e-06, 3.0740712e-06, -5.9681433e-06, 3.626002e-06, 7.487342e-06, 2.642479e-06, -4.9963705e-06, -5.125625e-06, 1.8369469e-06, 2.2820116e-07, -1.2280577e-06, 3.7428165e-07, 2.554675e-06, 3.9706126e-07]
mushroom_ir = [4.373888e-06, -4.9138384e-06, 8.0840555e-06, -6.7274595e-06, 4.094613e-06, -7.5424255e-06, 9.182771e-06, -3.8974426e-06, 9.911429e-06, -1.3705036e-05, 3.422383e-05, -3.839868e-05, -0.00019172326, 0.00011527779, 0.0001520886, 0.00012412298, -0.00013692162, -0.00013208372, -4.3006567e-06, 0.00011475379, 5.5517354e-05, -7.088181e-05, -1.4412982e-05, -1.7430219e-05, 3.554604e-05, -1.7222818e-06, 5.6685712e-06, -1.4786573e-05, -8.356573e-06, -2.2435363e-05, 6.9735893e-06, 1.7591914e-06, 7.1914546e-06, -2.301219e-06, 6.038217e-06, -5.537905e-06, -7.309603e-06, -4.380904e-06, 8.4743015e-06, 6.69175e-06, -4.441023e-06, -5.500982e-07, 8.955039e-06, 1.6854503e-06, -4.1729963e-06, -1.1186148e-06, 3.3305216e-06, -2.0207767e-06, -6.554812e-06, -2.0501457e-06, 9.365622e-06, 5.435293e-07, -1.078435e-06, 2.5975182e-06, 3.1110858e-06, -1.3232268e-06, -4.4947487e-06, 1.7805148e-06, 3.765102e-06, -3.5659642e-07, -4.689883e-06, 5.412453e-06, 1.2596523e-06, 1.0665299e-06]

fb_scale = np.power(10, 50 / 20)

tulip_ir = np.array(tulip_ir) * fb_scale
mushroom_ir = np.roll(np.array(mushroom_ir), 8) * fb_scale

IR = tulip_ir

N = 4
M = 32 * 1

w = np.concatenate((IR, np.zeros(N * M - len(IR))))
W = np.fft.rfft(w.reshape((-1, M)), 2*M, axis=1)
zi = np.zeros(len(IR) - 1)
WW = np.abs(W) ** 2

A = 0.999999
P_initial = np.max(WW) / 2
partial_constrain = True

ft = PFDKF(N, M, A, P_initial, partial_constrain)

# ft = FDKF(M * N, 0.99, P_initial)


d, fs = sf.read('in.wav')
# d, fs = sf.read('music.wav')

num_block = len(d) // M
mix = np.zeros((num_block * M, 8))
audio = np.zeros((num_block * M, 6))


out_buf = ring_buffer(fs, -M)
fb_buf = ring_buffer(fs, -M)


gain = 30 if len(sys.argv) < 2 else int(sys.argv[1])
scale = np.power(10, gain / 20)

import random
np.random.seed(random.SystemRandom().randint(1, 1024))

pulse_n = 20 * fs // M

noise = np.random.uniform(-1., 1., len(IR))  / 10

P = []
H = []
leak = []

for n in range(num_block):
    noise_n = np.random.uniform(-1., 1., M) / (1 << 28)
    x_n = out_buf.get(M)
    fb_n = fb_buf.get(M)
    near_n = d[n*M:(n+1)*M] + noise_n
    # near_n = d[n*M:(n+1)*M]
    d_n = near_n + fb_n
    # d_n = d[n*M:(n+1)*M]

    if 1:
        e_n = ft.filt(x_n, d_n)
        ft.update(e_n)
    else:
        e_n = d_n

    if n == 200:
        ft.freeze = 1

    if n == pulse_n:
        IR = mushroom_ir
        w = np.concatenate((IR, np.zeros(N * M - len(IR))))
        W = np.fft.rfft(w.reshape((-1, M)), 2*M, axis=1)

        ft.freeze = 0
    
    if n == (2 * pulse_n - 1000):
        scale = np.power(10, 40 / 20)
        # ft.P = np.abs(ft.H) ** 2 * 0.05

    out_n = np.clip(e_n * scale, -1, 1)

    out_buf.put(out_n)

    if (n & 0xFF) == 0:
        noise = np.random.uniform(-1., 1., len(IR))  / 20
    noise = 0

    fb, zi = signal.lfilter(IR * (1 + noise), 1, out_n, zi=zi)
    fb_buf.put(fb)

    r1 = n * M
    r2 = r1 + M
    mix[r1:r2, 0] = out_n
    mix[r1:r2, 1] = (e_n - near_n) * scale
    mix[r1:r2, 2] = np.sum(ft.P) * 100
    mix[r1:r2, 3] = np.sum(np.abs(W - ft.H)**2) * 100 * 2
    # mix[r1:r2, 4] = ft.y * scale
    mix[r1:r2, 5] = RMS2(mix[r1:r2, 1]) / (RMS2(mix[r1:r2, 4]) + 1e-10)
    # mix[r1:r2, 6] = np.sum(ft.leak)
    # mix[r1:r2, 7] = np.sum(np.abs(ft.dH) ** 2) * 10000

    audio[r1:r2, 0] = near_n
    audio[r1:r2, 1] = fb_n
    audio[r1:r2, 2] = d_n
    audio[r1:r2, 3] = e_n
    audio[r1:r2, 4] = e_n - near_n
    audio[r1:r2, 5] = out_n

    P.append(ft.P.copy())
    H.append(ft.H.copy())
    # leak.append(ft.leak)


mix = np.clip(mix, -1, 1)

sf.write('out.wav', mix[:, 0], fs, 'PCM_16')
sf.write('mix.wav', mix, fs, 'PCM_16')
# sf.write('audio.wav', audio, fs, 'PCM_16')

np.save('p.npy', P)
np.save('h.npy', H)
# np.save('leak.npy', leak)

sys.exit(0)

import matplotlib.pyplot as plt
plt.subplot(3, 1, 1)
plt.plot(IR[0] * fb_scale, label='IR0')
plt.plot(IR[1] * fb_scale, label='IR1')
plt.plot(h0, label='h0')
plt.plot(ft.h, label='h')
# plt.plot(h1, label='h1')

plt.legend()

plt.subplot(3, 1, 2)

plt.plot(IR[1] * fb_scale - ft.h[:len(IR[0])], label='IR1 - h')
plt.plot(IR[0] * fb_scale - h0[:len(IR[0])], label='IR0 - h0')
# plt.plot(IR * fb_scale - h1[:len(IR)], label='d1')
plt.legend()


# plt.subplot(3, 1, 3)
# plt.plot(p_data[100:])

plt.show()